package fun.easycode.webterminal.service.impl;


import cn.hutool.extra.ssh.JschUtil;
import com.jcraft.jsch.JSchException;
import fun.easycode.webterminal.common.contract.DynamicOperate;
import fun.easycode.webterminal.common.contract.ErrorCode;
import fun.easycode.webterminal.common.contract.PageDTO;
import fun.easycode.webterminal.common.contract.R;
import fun.easycode.webterminal.common.datafill.DataFill;
import fun.easycode.webterminal.dao.*;
import fun.easycode.webterminal.service.MachineTunnelService;
import fun.easycode.webterminal.service.cmd.MachineTunnelPageQry;
import fun.easycode.webterminal.service.cmd.MachineTunnelSaveCmd;
import fun.easycode.webterminal.service.cmd.MachineTunnelSaveCmdAssembler;
import fun.easycode.webterminal.service.dto.MachineTunnelDTO;
import fun.easycode.webterminal.service.dto.MachineTunnelDTOAssembler;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.sshd.client.SshClient;
import org.apache.sshd.client.keyverifier.AcceptAllServerKeyVerifier;
import org.apache.sshd.client.session.ClientSession;
import org.apache.sshd.common.forward.PortForwardingEventListener;
import org.apache.sshd.common.session.Session;
import org.apache.sshd.common.util.net.SshdSocketAddress;
import org.apache.sshd.server.forward.AcceptAllForwardingFilter;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

@Slf4j
@Service
public class MachineTunnelServiceImpl implements MachineTunnelService {

    @Resource
    private MachineTunnelRepository machineTunnelRepository;
    @Resource
    private MachineRepository machineRepository;

    private final SshClient client;

    // 线程安全map
    private final Map<String , ClientSession> sessionMap = new ConcurrentHashMap<>();
    private final Map<String, com.jcraft.jsch.Session> jschSessionMap = new ConcurrentHashMap<>();
    private final Map<String, String> startMap = new ConcurrentHashMap<>();
    // 使用ConcurrentHashMap来存储每个ID对应的锁对象
    private final ConcurrentHashMap<String, Lock> locks = new ConcurrentHashMap<>();

    // 创建调度任务线程池
    private final ScheduledExecutorService scheduledExecutorService = new ScheduledThreadPoolExecutor(2);

    public MachineTunnelServiceImpl(){
        this.client = SshClient.setUpDefaultClient();
        this.client.setForwardingFilter(AcceptAllForwardingFilter.INSTANCE);
        this.client.setServerKeyVerifier(AcceptAllServerKeyVerifier.INSTANCE);
        this.client.start();
    }

    @Override
    public R<String> delete(String id) {
        boolean r = machineTunnelRepository.removeById(id);
        return r? R.message("删除成功") : R.error("删除失败");
    }

    @Override
    public R<MachineTunnelDTO> getById(String id) {

        MachineTunnel machineTunnel = machineTunnelRepository.getById(id);

        if(machineTunnel == null){
            return R.error("未找到该隧道");
        }

        return R.success(MachineTunnelDTOAssembler.INSTANCE.to(machineTunnel));
    }

    @Override
    @DataFill
    public R<PageDTO<MachineTunnelDTO>> page(MachineTunnelPageQry cmd) {
        cmd.validate();

        PageDTO<MachineTunnelDTO> page = DynamicOperate.page(cmd, machineTunnelRepository.getBaseMapper(), MachineTunnelDTOAssembler.INSTANCE::to);


//        List<String> machineIdList = page.getData().stream().map(MachineTunnelDTO::getMachineId).toList();
//
//        List<Machine> machines = machineRepository.listByIds(machineIdList);
//
//        Map<String, Machine> machineMap = machines.stream().collect(Collectors.toMap(Machine::getId, Function.identity()));
//
        page.getData().forEach(machineTunnelDTO -> {
            // 填充状态
            machineTunnelDTO.setStart(startMap.containsKey(machineTunnelDTO.getId()));
        });

        return R.success(page);
    }

    @Override
    public R<String> save(MachineTunnelSaveCmd cmd) {
        cmd.validate();

        MachineTunnel tunnel = MachineTunnelSaveCmdAssembler.INSTANCE.to(cmd);

        boolean r = machineTunnelRepository.saveOrUpdate(tunnel);

        return r ? R.success("保存成功!") : R.error("保存失败!");
    }

    @Override
    public R<String> start(String id) {

        MachineTunnel machineTunnel = machineTunnelRepository.getById(id);

        if(machineTunnel == null){
            return R.error("未找到该隧道");
        }

        Machine machine = machineRepository.getById(machineTunnel.getMachineId());

        if(machine == null){
            return R.error("未找到该机器");
        }

        if(Objects.equals(startMap.getOrDefault(machineTunnel.getId(), "0"), "1")){
            return R.message("该隧道已启动");
        }

        R<String> r = null;

        Lock lock = getLock(id);

        log.info("开始锁");
        lock.lock();
        log.info("获取锁");
        try {
            if(Objects.equals(startMap.getOrDefault(machineTunnel.getId(), "0"), "1")){
                r = R.message("该隧道已启动");
            }else {
                if (machineTunnel.getType() == TunnelType.DYNAMIC) {
                    r = startDynamicTunnel(machineTunnel, machine);
                } else if (machineTunnel.getType() == TunnelType.LOCAL) {
                    r = startLocalTunnel(machineTunnel, machine);
                } else if (machineTunnel.getType() == TunnelType.REMOTE) {
                    r = startRemoteTunnel(machineTunnel, machine);
                }
            }
        }finally {
            log.info("释放锁");
            lock.unlock();
        }

        if(r != null){

            if(r.getCode() == ErrorCode.Success){
                startMap.put(machineTunnel.getId(), "1");
            }

            return r;
        }

        return R.error("未知隧道类型");
    }

    @Override
    public R<String> stop(String id) {
        MachineTunnel machineTunnel = machineTunnelRepository.getById(id);

        if(machineTunnel == null){
            return R.error("未找到该隧道");
        }

        Machine machine = machineRepository.getById(machineTunnel.getMachineId());

        if(machine == null){
            return R.error("未找到该机器");
        }

        R<String> r = null;

        Lock lock = getLock(id);
        lock.lock();
        try {
            if (machineTunnel.getType() == TunnelType.DYNAMIC) {
                r = stopDynamicTunnel(machineTunnel, machine);
            } else if (machineTunnel.getType() == TunnelType.LOCAL) {
                r = stopLocalTunnel(machineTunnel, machine);
            } else if (machineTunnel.getType() == TunnelType.REMOTE) {
                r = stopRemoteTunnel(machineTunnel, machine);
            }
        }finally {
            lock.unlock();
        }

        if(r != null){

            if(r.getCode() == ErrorCode.Success){
                startMap.remove(machineTunnel.getId());
            }

            return r;
        }

        return R.error("未知隧道类型");

    }

    @SneakyThrows
    private R<String> stopDynamicTunnel(MachineTunnel machineTunnel, Machine machine ) {
        ClientSession clientSession = sessionMap.get(machineTunnel.getId());

        clientSession.close();

        sessionMap.remove(machineTunnel.getId());

        return R.message("隧道停止成功");
    }

    private R<String> stopLocalTunnel(MachineTunnel machineTunnel, Machine machine ) {
        com.jcraft.jsch.Session session = jschSessionMap.get(machineTunnel.getId());
        if(session != null){
            if(session.isConnected()){
                JschUtil.unBindPort(session, machineTunnel.getLocalPort());
            }
            session.disconnect();
            jschSessionMap.remove(machineTunnel.getId());
        }
        return R.message("隧道停止成功");
    }

    private R<String> stopRemoteTunnel(MachineTunnel machineTunnel, Machine machine) {
        com.jcraft.jsch.Session session = jschSessionMap.get(machineTunnel.getId());
        if(session != null){
            if(session.isConnected()){
                try {
                    session.delPortForwardingR(machineTunnel.getRemotePort());
                } catch (JSchException e) {
                    throw new RuntimeException(e);
                }
//                JschUtil.unBindPort(session, machineTunnel.getRemotePort());
            }
            session.disconnect();
            jschSessionMap.remove(machineTunnel.getId());
        }
        return R.message("隧道停止成功");
    }

    private R<String> startLocalTunnel(MachineTunnel machineTunnel, Machine machine ) {

        com.jcraft.jsch.Session session = JschUtil.openSession(machine.getIp(), machine.getPort(), machine.getUsername(), machine.getPassword());

        jschSessionMap.put(machineTunnel.getId(), session);

        boolean r = JschUtil.bindPort(session, machineTunnel.getRemoteIp(), machineTunnel.getRemotePort(), machineTunnel.getLocalIp(), machineTunnel.getLocalPort());

        return r? R.message("隧道启动成功") : R.error("隧道启动失败");
    }

    private R<String> startRemoteTunnel(MachineTunnel machineTunnel, Machine machine) {
        com.jcraft.jsch.Session session = JschUtil.openSession(machine.getIp(), machine.getPort(), machine.getUsername(), machine.getPassword());

        jschSessionMap.put(machineTunnel.getId(), session);

        boolean r = JschUtil.bindRemotePort(session, machineTunnel.getRemotePort(), machineTunnel.getLocalIp(), machineTunnel.getLocalPort());
        return r? R.message("隧道启动成功") : R.error("隧道启动失败");
    }

    @SneakyThrows
    private R<String> startDynamicTunnel(MachineTunnel machineTunnel, Machine machine) {

        if(sessionMap.get(machineTunnel.getId()) != null){
            return R.error("该机器已存在隧道");
        }

        // using the client for multiple sessions...
        ClientSession session = client.connect(machine.getUsername(), machine.getIp(), machine.getPort()).verify()
                .getSession();

        sessionMap.put(machineTunnel.getId(), session);

        // IF you use password to login provide here
        // session.addPasswordIdentity(BASTION_SERVER_PASSWORD); // for password-based
        // authentication

        session.addPasswordIdentity(machine.getPassword());
        // authentication
        // Note: can add BOTH password AND public key identities - depends on the
        // client/server security setup
        session.auth().verify(10000);
        // start using the session to run commands, do SCP/SFTP, create local/remote
        // port forwarding, etc...
        session.addPortForwardingEventListener(new PortForwardingEventListener() {
            @Override
            public void establishedDynamicTunnel(Session session, SshdSocketAddress local,
                                                 SshdSocketAddress boundAddress, Throwable reason) throws IOException {
                // TODO Auto-generated method stub
                PortForwardingEventListener.super.establishedDynamicTunnel(session, local, boundAddress, reason);
                log.info("Dynamic forward Tunnel is Ready host {} port {} .", machineTunnel.getLocalIp(), machineTunnel.getLocalPort());
            }
        });

        session.startDynamicPortForwarding(new SshdSocketAddress(machineTunnel.getLocalIp(), machineTunnel.getLocalPort()));

        return R.message("启动成功！");
    }

    @PostConstruct
    public void scheduled(){
        log.info("scheduled start");
        // 处理socket5转发
        scheduledExecutorService.scheduleWithFixedDelay(() -> {
            sessionMap.forEach((key, value) -> {
                Lock lock = getLock(key);
                lock.lock();
                try {
                    if (!value.isOpen()) {
                        log.info("scheduled session {} is closed", key);
                        stop(key);
                        log.info("scheduled session {} is start", key);
                        R<String> start = start(key);
                        if (start.getCode() != ErrorCode.Success) {
                            log.error("scheduled session {} start error msg {}", key, start.getMessage());
                        }
                    }
                }finally {
                    lock.unlock();
                }
            });
        }, 1, 2, TimeUnit.SECONDS);

        // 处理jsch端口转发
        scheduledExecutorService.scheduleWithFixedDelay(() -> {
            jschSessionMap.forEach((key, value) -> {
                Lock lock = getLock(key);
                lock.lock();
                try {
                    if (!value.isConnected()) {
                        log.info("scheduled jschSession {} is closed", key);
                        stop(key);
                        log.info("scheduled jschSession {} is start", key);
                        R<String> start = start(key);
                        if (start.getCode() != ErrorCode.Success) {
                            log.error("scheduled jschSession {} start error msg {}", key, start.getMessage());
                        }
                    }
                }finally {
                    lock.unlock();
                }
            });
        }, 1, 2, TimeUnit.SECONDS);
    }

    /**
     * 获取指定ID的锁
     *
     * @param id 需要加锁的ID
     * @return 返回对应ID的锁对象
     */
    public Lock getLock(String id) {
        // 使用computeIfAbsent方法来确保每个ID只有一个锁对象
        return locks.computeIfAbsent(id, k -> new ReentrantLock());
    }
}
